package cn.aylog.simpleTree;

import java.io.*;
import java.util.*;

/**
 * 赫夫曼树（wpl： weight path length）最小的树
 * 构建规则：
 * 1、每一个元素都是一个最小二叉树
 * 2、取出2个最小的树，构建成一棵新的二叉树
 * 3、将2个最小的元素删除，将新树加入到原数组
 */
public class HuffmanCodeDemo {
    public static void main(String[] args) {
        /*
        // 赫夫曼编码
        String str = "i use use use use java a java application to do do do something";
        byte[] bytes = str.getBytes();
        // System.out.println("src:"+Arrays.toString(bytes));
        // List<HuffmanCodeNode> nodeList = countWeight(bytes);
        // HuffmanCodeNode huffmanTree = createHuffmanTree(nodeList);
        // // 获取赫夫曼编码表
        // Map<Byte, String> huffmanCode = createHuffmanCode(huffmanTree);
        Map<Byte, String> huffmanCode = createHuffmanCode(bytes);
        // for(Map.Entry<Byte, String> entry : huffmanCode.entrySet()) {
        //     System.out.println(entry.getKey() + "->" + entry.getValue());
        // }
        // 压缩
        byte[] zip = zip(bytes, huffmanCode);
        // 解压
        byte[] unzip = unzip(zip, huffmanCode);
        // System.out.println("dis:"+Arrays.toString(unzip));
        for (byte b : unzip) {
            System.out.printf("%c", b);
        }
        System.out.println();
        */

        String srcPath = "d:/data/yarn.lock";
        String dstPath = "d:/data/yarn.zip";
        // zipFile(srcPath, dstPath);
        // System.out.println("压缩成功~");
        unzipFile(dstPath, srcPath+"1");
        System.out.println("解压成功~");

    }

    public static void unzipFile(String src, String dst) {
        File srcFile = new File(src);
        File dstFile = new File(dst);
        ObjectInputStream ois = null;
        BufferedOutputStream bos = null;
        try {
            ois = new ObjectInputStream(new FileInputStream(srcFile));
            byte[] bytes = (byte[])ois.readObject();
            Map<Byte, String> huffmanCode = (Map<Byte, String>)ois.readObject();
            byte[] unzipBytes = unzip(bytes, huffmanCode);
            bos = new BufferedOutputStream(new FileOutputStream(dstFile));
            bos.write(unzipBytes);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                bos.close();
                ois.close();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    public static void zipFile(String src, String dst) {
        File srcFile = new File(src);
        File dstFile = new File(dst);
        BufferedInputStream bis = null;
        ObjectOutputStream oos = null;
        byte[] bytes = new byte[(int)srcFile.length()];
        try {
            bis = new BufferedInputStream(new FileInputStream(srcFile));
            bis.read(bytes);
            Map<Byte, String> huffmanCode = createHuffmanCode(bytes);
            byte[] zipBytes = zip(bytes, huffmanCode);
            oos = new ObjectOutputStream(new FileOutputStream(dstFile));
            oos.writeObject(zipBytes);
            oos.writeObject(huffmanCode);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                oos.close();
                bis.close();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    public static Map<Byte, String> createHuffmanCode(byte[] bytes) {
        // 计算权重,将数组转成赫夫曼节点
        List<HuffmanCodeNode> nodeList = countWeight(bytes);
        // 生成赫夫曼树
        HuffmanCodeNode huffmanTree = createHuffmanTree(nodeList);
        // 生成赫夫曼节点
        Map<Byte, String> huffmanCode = createHuffmanCode(huffmanTree);
        return huffmanCode;
    }

    /**
     * 解码
     *
     * @param bytes       待解码的数据
     * @param huffmanCode 编码表
     */
    public static byte[] unzip(byte[] bytes, Map<Byte, String> huffmanCode) {
        // 建立map映射,便于查找
        Map<String, Byte> unzipCode = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffmanCode.entrySet()) {
            unzipCode.put(entry.getValue(), entry.getKey());
        }
        StringBuilder keyBuilder = new StringBuilder(); // 用于拼接key
        StringBuilder temp = new StringBuilder(); // 临时存储
        List<Byte> byteList = new ArrayList<>(); // 保存转换后的数据
        int i = 0, len = bytes.length - 1, flagByte = bytes[len];
        // 常规处理
        for (int x = 0; x < len - 1; x++) {
            temp.append(byteToBit(bytes[x]));
            for (i = 0; i < 8; i++) {
                // 将每个字符拼接起来,根据映射表还原数据
                keyBuilder.append(temp.charAt(i));
                if (unzipCode.containsKey(keyBuilder.toString())) {
                    byteList.add(unzipCode.get(keyBuilder.toString()));
                    // 清空
                    keyBuilder.delete(0, keyBuilder.length());
                }
            }
            temp.delete(0, temp.length());
        }
        // 处理最后一位
        if (flagByte == 1) {
            temp.append(Integer.toBinaryString(bytes[len - 1]));
        } else {
            temp.append(byteToBit(bytes[len - 1]));
        }
        for (i = 0; i < temp.length(); i++) {
            // 将每个字符拼接起来,根据映射表还原数据
            keyBuilder.append(temp.charAt(i));
            if (unzipCode.containsKey(keyBuilder.toString())) {
                byteList.add(unzipCode.get(keyBuilder.toString()));
                // 清空
                keyBuilder.delete(0, keyBuilder.length());
            }
        }
        temp.delete(0, temp.length());
        byte[] resBytes = new byte[byteList.size()];
        i = 0;
        for (byte b : byteList) {
            resBytes[i++] = b;
        }
        return resBytes;
    }

    /**
     * 开始编码
     *
     * @param bytes       待编码的数据
     * @param huffmanCode 赫夫曼编码
     */
    public static byte[] zip(byte[] bytes, Map<Byte, String> huffmanCode) {
        StringBuilder resBuilder = new StringBuilder(); // 最后结果
        for (byte b : bytes) {
            // 将转换后的字节存储到临时对象
            resBuilder.append(huffmanCode.get(b));
        }
        int l = resBuilder.length();
        int len = l / 8;
        if (l % 8 != 0) {
            len += 1;
        }
        len += 1; // 用于存储标志位
        byte[] resBytes = new byte[len];
        int index = 0;
        for (int i = 0; i < l; i += 8) {
            if (i + 8 < l) {
                resBytes[index++] = (byte) (Integer.parseInt(resBuilder.substring(i, i + 8), 2) - 256);
            } else {
                String temp = resBuilder.substring(i);
                if (temp.charAt(0) == '0') {
                    // 0-低位补0
                    int fix0 = 8 - temp.length(); // 计算需要补几位0
                    for (int x = 0; x < fix0; x++) {
                        temp = temp + "0";
                    }
                    resBytes[index++] = (byte) (Integer.parseInt(temp, 2) - 256);
                    resBytes[index] = 0;
                } else {
                    // 1-高位补0
                    resBytes[index++] = (byte) (Integer.parseInt(temp, 2) - 256);
                    resBytes[index] = 1;
                }
            }
        }
        return resBytes;
    }

    /**
     * 把byte转为字符串的bit
     */
    public static String byteToBit(byte b) {
        return ""
                + (byte) ((b >> 7) & 0x1) + (byte) ((b >> 6) & 0x1)
                + (byte) ((b >> 5) & 0x1) + (byte) ((b >> 4) & 0x1)
                + (byte) ((b >> 3) & 0x1) + (byte) ((b >> 2) & 0x1)
                + (byte) ((b >> 1) & 0x1) + (byte) ((b >> 0) & 0x1);
    }


    /**
     * 从赫夫曼树中构建可变字长编码（vlc）
     * 左边为0， 右边为1
     *
     * @param root 赫夫曼树的根节点
     */
    public static Map<Byte, String> createHuffmanCode(HuffmanCodeNode root) {
        StringBuilder builder = new StringBuilder();
        Map<Byte, String> huffmanCode = new HashMap<>();
        getCode(root, "", builder, huffmanCode);
        return huffmanCode;
    }

    /**
     * 递归生成编码
     *
     * @param node        某个节点
     * @param ch          编码
     * @param builder     字符构建
     * @param huffmanCode 存储赫夫曼编码
     */
    private static void getCode(HuffmanCodeNode node, String ch, StringBuilder builder, Map<Byte, String> huffmanCode) {
        builder.append(ch);
        if (node.getLeft() != null) getCode(node.getLeft(), "0", builder, huffmanCode);
        if (node.getRight() != null) getCode(node.getRight(), "1", builder, huffmanCode);
        if (node.getValue() != null) {
            huffmanCode.put(node.getValue().byteValue(), builder.toString());
        }
        // 实际有效的节点必定是叶子节点，所以打印完之后返回上一层节点，必须删除当前节点的标识
        if (builder.length() > 0) {
            builder.deleteCharAt(builder.length() - 1);
        }
    }

    /**
     * 统计byte数组的每个字符的权重
     *
     * @param bytes 待统计的字符数组
     */
    public static List<HuffmanCodeNode> countWeight(byte[] bytes) {
        Map<Byte, Integer> counts = new HashMap<>();
        for(byte b : bytes) {
            if(counts.containsKey(b)) {
                counts.put(b, counts.get(b)+1);
            } else {
                counts.put(b, 1);
            }
        }
        List<HuffmanCodeNode> nodes = new ArrayList<>();
        // 权重统计
        for(Map.Entry<Byte, Integer> entry:counts.entrySet()) {
            nodes.add(new HuffmanCodeNode(entry.getKey().intValue(), entry.getValue()));
        }
        return nodes;
    }

    /**
     * 将给定的数组构建成一个赫夫曼树
     *
     * @return 转换后的赫夫曼树的根节点
     */
    public static HuffmanCodeNode createHuffmanTree(List<HuffmanCodeNode> nodes) {
        // 数组转List<Node>
        HuffmanCodeNode leftNode, rightNode, parentNode;
        while (nodes.size() > 1) {
            //排序
            Collections.sort(nodes);
            //构建树
            leftNode = nodes.get(0);
            rightNode = nodes.get(1);
            parentNode = new HuffmanCodeNode(null, leftNode.getWeight() + rightNode.getWeight());
            parentNode.setLeft(leftNode);
            parentNode.setRight(rightNode);
            // 从原数组去除
            nodes.remove(leftNode);
            nodes.remove(rightNode);
            // 新增节点
            nodes.add(parentNode);
        }
        return nodes.get(0);
    }
}

class HuffmanCodeNode implements Comparable<HuffmanCodeNode> {
    private Integer value;
    private Integer weight; // 权重
    private HuffmanCodeNode left;
    private HuffmanCodeNode right;

    public HuffmanCodeNode(Integer value, Integer weight) {
        this.value = value;
        this.weight = weight;
    }

    public Integer getValue() {
        return value;
    }

    public void setValue(Integer value) {
        this.value = value;
    }

    public Integer getWeight() {
        return weight;
    }

    public void setWeight(Integer weight) {
        this.weight = weight;
    }

    public HuffmanCodeNode getLeft() {
        return left;
    }

    public void setLeft(HuffmanCodeNode left) {
        this.left = left;
    }

    public HuffmanCodeNode getRight() {
        return right;
    }

    public void setRight(HuffmanCodeNode right) {
        this.right = right;
    }

    @Override
    public int compareTo(HuffmanCodeNode o) {
        return this.getWeight() - o.getWeight();
    }

    @Override
    public String toString() {
        return "HuffmanCodeNode{" +
                "value=" + value +
                ", weight=" + weight +
                '}';
    }
}
